from langchain_core.output_parsers import JsonOutputParser
from langchain.agents import initialize_agent, AgentType
from pydantic import BaseModel, Field
from app.bailian.common import create_calc_tools, llm, chat_prompt_template

class Output(BaseModel):
    args: str = Field(description="工具入参")
    result: str = Field(description="工具计算结果")
    thinking: str = Field(description="思考过程")

parser = JsonOutputParser(pydantic_object = Output)
format_instructions = parser.get_format_instructions()

#print(format_instructions);

agent = initialize_agent(
    tools=create_calc_tools(),
    llm=llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,  # 是否打印中间思考过程
)

prompt = chat_prompt_template.format_messages(
    role="计算",
    domain="使用工具进行数学计算",
    question=f"100+100等于多少？请严格按照以下格式输出结果：\n{format_instructions}"
)

response = agent.invoke(prompt)

print(type(response))
print(response)
print(type(response["output"]))
print(response["output"])